package org.gwh.airagknowledge.core.embedding;

import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.gwh.airagknowledge.entity.DocumentChunk;
import org.gwh.airagknowledge.entity.VectorEmbedding;
import org.gwh.airagknowledge.repository.VectorEmbeddingRepository;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import java.nio.ByteBuffer;
import java.util.List;

@Slf4j
@Component
@RequiredArgsConstructor
public class EmbeddingGenerator {

    private final VectorEmbeddingRepository vectorEmbeddingRepository;
    private final EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

    @Transactional
    public VectorEmbedding generateEmbedding(DocumentChunk chunk) {
        try {
            // 生成嵌入向量
            float[] vector = embeddingModel.embed(chunk.getContent()).content().vector();
            
            // 转换为字节数组
            byte[] vectorBytes = floatArrayToByteArray(vector);
            
            // 创建并保存向量嵌入
            VectorEmbedding embedding = VectorEmbedding.builder()
                    .chunk(chunk)
                    .embedding(vectorBytes)
                    .build();
            
            return vectorEmbeddingRepository.save(embedding);
        } catch (Exception e) {
            log.error("Error generating embedding for chunk ID: {}", chunk.getId(), e);
            throw new RuntimeException("Failed to generate embedding", e);
        }
    }

    public float[] generateQueryEmbedding(String query) {
        return embeddingModel.embed(query).content().vector();
    }
    
    public byte[] generateQueryBytesVector(String query) {
        return floatArrayToByteArray(generateQueryEmbedding(query));
    }
    
    public List<VectorEmbedding> findSimilarEmbeddings(Long knowledgeBaseId, String query, int limit) {
        byte[] queryVector = generateQueryBytesVector(query);
        return vectorEmbeddingRepository.findSimilarEmbeddings(knowledgeBaseId, queryVector, limit);
    }
    
    // 工具方法：浮点数组转字节数组
    private byte[] floatArrayToByteArray(float[] floats) {
        ByteBuffer buffer = ByteBuffer.allocate(floats.length * 4); // 每个float占4字节
        for (float value : floats) {
            buffer.putFloat(value);
        }
        return buffer.array();
    }
    
    // 工具方法：字节数组转浮点数组
    public static float[] byteArrayToFloatArray(byte[] bytes) {
        ByteBuffer buffer = ByteBuffer.wrap(bytes);
        float[] floats = new float[bytes.length / 4]; // 每个float占4字节
        for (int i = 0; i < floats.length; i++) {
            floats[i] = buffer.getFloat();
        }
        return floats;
    }
} 